﻿using System;
using System.Collections.Generic;
using System.Diagnostics.Metrics;
using System.Linq;
using System.Text;
using System.Threading.Channels;
using System.Threading.Tasks;

namespace JobDemo01.Job
{
    public interface IJobQueue
    {
        IJob Current { get; }

        bool IsBusy { get; }

        ISubscription<IJob> Subscribe();

        IJob UnsafeRun(JobPlan plan, JobRunOptions options);

        Task<WaferNavigatorOptions> TestAlignment(WaferMap map, Pattern pattern, JobEnvironment env);

        Task<WaferNavigatorOptions> TestAlignment(WaferMap map, List<Pattern> patterns, JobEnvironment env);

        Task TestMeasurement(JobPlan plan, JobRunOptions options);

        // FIX ME
        Task MoveTo(JobPlan plan, int scanIndex)
        {
            return MoveTo(plan.ScanPoints[scanIndex]);
        }

        Task MoveTo(ScanPoint point);

        Task<Grid> MakeGrid(JobPlan plan, JobRunOptions options);
    }

    public class JobQueue(FrameDispatcher dispatcher,
        IStageService stageService,
        IImagingService imagingService,
        ISEMClient semClient,
        JobResultRepo repo) : IJobQueue
    {
        readonly ILogger logger = Log.ForContext<JobQueue>();

        readonly object lockObj = new();
        readonly MessageQueue<IJob> queue = new();

        bool isBusy;

        IJob current;
        MessageQueue<PatternResult> resultQueue;

        public IJob Current
        {
            get
            {
                lock (lockObj)
                {
                    return current;
                }
            }
        }

        public bool IsBusy
        {
            get
            {
                lock (lockObj)
                {
                    return isBusy;
                }
            }
        }

        // assume we have the lock
        void SetCurrent(IJob job, MessageQueue<PatternResult> q)
        {
            resultQueue?.Dispose();

            current = job;
            resultQueue = q;
            isBusy = true;
            queue.Send(job);
        }

        Job NewJob(WaferMap map)
        {
            lock (lockObj)
            {
                if (isBusy)
                {
                    throw new InvalidOperationException("Job queue is busy.");
                }

                var ctx = new JobContext(this, map);
                var job = new Job(ctx);
                SetCurrent(job, ctx.MessageQueue);
                return job;
            }
        }

        public ISubscription<IJob> Subscribe()
        {
            return queue.Subscribe();
        }

        public IJob UnsafeRun(JobPlan plan, JobRunOptions options)
        {
            lock (lockObj)
            {
                if (isBusy)
                {
                    throw new InvalidOperationException("Job queue is busy.");
                }

                var ctx = new JobContext(this, plan.WaferMap);
                var planTask = MakeExecutionPlan(ctx.Id, plan, options);
                _ = Run(new ImagingContext(ctx), planTask, RunImagingPipeline, RunPostProcessing);
                var job = new ConfirmationJob(ctx, planTask, Retry);
                SetCurrent(job, ctx.MessageQueue);
                return job;
            }
        }

        public Task<WaferNavigatorOptions> TestAlignment(WaferMap map, Pattern pattern, JobEnvironment env)
        {
            var job = NewJob(map);
            var plan = Task.FromResult(new ExecutionPlan
            {
                Environment = env,
                Options = new()
            });
            return Run(job.ImagingContext, plan, async (ctx, plan) =>
            {
                FixHorizontalAlignment(pattern.Coordinates);

                var nav = stageService.WaferNavigator(stageService.Options with
                {
                    Map = map.WaferMapSpec()
                });
                var pipeline = ImagingPipeline(ctx, nav, plan);
                var alignment = new WaferAlignment(imagingService, stageService, pipeline);
                var options = await alignment.Align(pattern).ConfigureAwait(false);
                await semClient.Settings.SetAngle(options.Angle);
                return options;
            });
        }

        public Task<WaferNavigatorOptions> TestAlignment(WaferMap map, List<Pattern> patterns, JobEnvironment env)
        {
            var job = NewJob(map);
            var plan = Task.FromResult(new ExecutionPlan
            {
                Map = map,
                Alignment = OptimizeAlignment(patterns),
                Environment = env,
                Options = new()
            });
            return Run(job.ImagingContext, plan, async (ctx, plan) =>
            {
                var nav = await Align(ctx, plan);
                return nav.Options;
            });
        }

        public Task TestMeasurement(JobPlan plan, JobRunOptions options)
        {
            lock (lockObj)
            {
                if (isBusy)
                {
                    throw new InvalidOperationException("Job queue is busy.");
                }

                var ctx = new JobContext(this, plan.WaferMap);
                var planTask = MakeExecutionPlan(ctx.Id, plan, options);
                var job = new ConfirmationJob(ctx, planTask, Retry);
                SetCurrent(job, ctx.MessageQueue);
                return Run(new ImagingContext(ctx), planTask, RunImagingPipeline, RunPostProcessing);
            }
        }

        public Task MoveTo(ScanPoint p)
        {
            var job = NewJob(null);
            var plan = Task.FromResult(new ExecutionPlan
            {
                Alignment = [],
                ScanPoints = [p],
                Environment = new(),
                Options = new(),
            });
            return Run(job.ImagingContext, plan, RunImagingPipeline, null);
        }

        public Task<Grid> MakeGrid(JobPlan plan, JobRunOptions options)
        {
            var job = NewJob(plan.WaferMap);
            var planTask = Task.FromResult(new ExecutionPlan
            {
                Environment = plan.Environment,
                Options = new(),
            });
            return Run(job.ImagingContext, planTask, async (ctx, plan) =>
            {
                var nav = await Align(ctx, plan).ConfigureAwait(false);
                var spec = nav.Options.Map;
                var w = spec.ChipW;
                var h = spec.ChipH;
                var grid = new Grid()
                {
                    X = spec.ArrayX,
                    Y = spec.ArrayY,
                    Width = w,
                    Height = h,
                    Coordinates = new double[spec.ArrayX * spec.ArrayY * 2],
                    Angle = nav.Options.Angle,
                };
                var pattern = plan.Alignment[1];
                var c = pattern.Coordinates[0];
                if (plan.ScanPoints.Count > 0)
                {
                    var p = plan.ScanPoints[0];
                    c = new Coordinate(p.X, p.Y, p.Dx, p.Dy, p.AddrOffsetX, p.AddrOffsetY);
                    pattern = p.Pattern;
                }

                // assume we own the plan
                pattern.Kind = 2;
                pattern.Reference.DoCentering = true;

                var pipeline = ImagingPipeline(ctx, nav, plan);
                for (var i = 0; i < spec.ArrayX; i++)
                {
                    for (var j = 0; j < spec.ArrayY; j++)
                    {
                        c = c with
                        {
                            X = i,
                            Y = j
                        };

                        if (!spec.IsInside(c.X, c.Y, c.Dx, c.Dy) ||
                            !spec.IsInside(c.X, c.Y, c.Dx + c.AddrOffsetX, c.Dy + c.AddrOffsetY))
                        {
                            continue;
                        }

                        var mp = await pipeline.MoveTo(c, pattern);
                        var (sx, sy) = await nav.MotionControl.GetCoordinate();

                        if (mp.Score < mp.Pattern.Reference.Threshold)
                        {
                            continue;
                        }

                        var cond = semClient.Condition;
                        var mat = Transformation.NewTransformation(cond.PixelSize, cond.ImageAngle);
                        var offset = mat.Transform(new Vector(-mp.MatchDx, -mp.MatchDy));
                        grid[i, j] = new Vector(sx, sy) + offset;
                    }
                }

                return grid;
            });
        }

        async Task<ExecutionPlan> MakeExecutionPlan(Guid id, JobPlan plan, JobRunOptions options)
        {
            using var ctx = repo.NewContext();
            var patterns = new List<PatternCopy>();
            foreach (var p in plan.Alignment)
            {
                patterns.Add(new PatternCopy
                {
                    Kind = p.Kind,
                    Ordinal = p.Ordinal,
                    Pattern = Copy(p)
                });
            }

            var scanOrder = 0;
            var coordinates = new List<ChipCoordinate>();
            var parameters = new Dictionary<(int, int), AdvVerifyModels.Entities.Parameter>();
            foreach (var s in plan.ScanPoints)
            {
                coordinates.Add(new ChipCoordinate(s.X, s.Y, s.Dx, s.Dy));

                var pattern = s.Pattern;
                if (!parameters.ContainsKey((s.PatternId, 0)))
                {
                    var copy = new PatternCopy
                    {
                        Kind = pattern.Kind,
                        Ordinal = pattern.Ordinal,
                        Pattern = Copy(s.Pattern)
                    };
                    patterns.Add(copy);

                    if (pattern.CDPoints.Count == 0)
                    {
                        parameters[(s.PatternId, 0)] = new AdvVerifyModels.Entities.Parameter
                        {
                            Name = $"P{s.Pattern.Ordinal}",
                            PatternCopy = copy,
                        };
                    }
                    else
                    {
                        for (var i = 0; i < pattern.CDPoints.Count; i++)
                        {
                            parameters[(s.PatternId, i)] = new AdvVerifyModels.Entities.Parameter
                            {
                                Name = pattern.CDPoints[i].Name,
                                Index = i,
                                PatternCopy = copy,
                            };
                        }
                    }
                }

                var n = pattern.CDPoints.Count == 0 ? 1 : pattern.CDPoints.Count;
                for (int i = 0; i < n; i++)
                {
                    var param = parameters[(s.PatternId, i)];
                    param.Measurements.Add(new Measurement
                    {
                        ScanIndex = scanOrder,
                        X = s.X,
                        Y = s.Y,
                        Status = MeasurementStatus.NotRun,
                    });
                }

                scanOrder += 1;
            }

            var result = new JobResult
            {
                Start = DateTimeOffset.Now.ToUnixTimeMilliseconds(),
                IsConfirmed = !plan.Environment.RequireConfirmation,
                Name = plan.Name,
                Operator = plan.Operator,
                LotId = plan.LotId,
                Process = plan.Process,
                WaferId = plan.WaferId,
                ConditionNo = plan.ConditionNo,
                RecipeId = plan.Id,
                OutputDirectory = options.OutputDirectory,
                Patterns = patterns,
                Parameters = [.. parameters.Values],
                JobId = id
            };
            ctx.Add(result);
            await ctx.SaveChangesAsync().ConfigureAwait(false);

            var mId = new Dictionary<(int, int), int>();
            foreach (var p in result.Parameters)
            {
                foreach (var m in p.Measurements)
                {
                    mId.Add((m.ScanIndex, p.Index), m.Id);
                }
            }

            return new ExecutionPlan
            {
                Id = result.Id,
                Measurements = mId,
                Map = plan.WaferMap,
                Environment = plan.Environment,
                Options = options,
                Alignment = OptimizeAlignment(plan.Alignment),
                ScanPoints = plan.ScanPoints,
                Estimation = new Estimation
                {
                    Time = coordinates.Count * 3,
                    Coordinates = coordinates
                }
            };
        }

        Task<T> Run<T>(ImagingContext ctx, Task<ExecutionPlan> planTask, Func<ImagingContext, ExecutionPlan, Task<T>> runImaging)
        {
            return Run(ctx, planTask, runImaging, NopPostProcessing);
        }

        async Task<T> Run<T>(ImagingContext ctx,
            Task<ExecutionPlan> planTask,
            Func<ImagingContext, ExecutionPlan, Task<T>> runImaging,
            Func<IAsyncEnumerable<PatternResult>, ExecutionPlan, Task> runProcessing)
        {
            try
            {
                var plan = await planTask.ConfigureAwait(false);
                var jobContext = ctx.JobContext;
                var pipeline = new ProcessingPipeline(jobContext.MessageQueue,
                    imagingService,
                    dispatcher,
                    plan.Environment,
                    plan.Options);
                var source = pipeline.Process(ctx.Source);
                var processing = runProcessing(source, plan);
                if (plan.Environment.RequireConfirmation)
                {
                    _ = RunProcessingWithContext(ctx, processing);
                }
                else
                {
                    _ = RunProcessingWithContext(ctx, jobContext.MessageQueue, processing);
                }

                var result = await runImaging(ctx, plan).ConfigureAwait(false);
                ctx.Complete(plan.Environment.RequireConfirmation);
                return result;
            }
            catch (Exception ex)
            {
                ctx.Complete(ex);
                throw;
            }
        }

        static async Task RunProcessingWithContext(ImagingContext ctx, Task processing)
        {
            try
            {
                await processing.ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                ctx.Complete(ex);
            }
        }

        static async Task RunProcessingWithContext(ImagingContext ctx, MessageQueue<PatternResult> queue, Task processing)
        {
            try
            {
                using (queue)
                {
                    await processing.ConfigureAwait(false);
                }
            }
            catch (Exception ex)
            {
                ctx.Complete(ex);
            }
        }

        IAsyncEnumerable<PatternResult> Retry(JobContext ctx, List<int> indices, Task<ExecutionPlan> planTask)
        {
            var channel = Channel.CreateUnbounded<PatternResult>();
            _ = Run(new ImagingContext(ctx, indices), planTask, RunImagingPipeline, HandleRetry);
            return channel.Reader.ReadAllAsync();

            async Task HandleRetry(IAsyncEnumerable<PatternResult> source, ExecutionPlan plan)
            {
                try
                {
                    using var writer = new ResultWriter(repo.NewContext(), plan);
                    await foreach (var result in source.ConfigureAwait(false))
                    {
                        await writer.Delete(result.Index).ConfigureAwait(false);
                        await writer.Write(result).ConfigureAwait(false);
                        channel.Writer.TryWrite(result);
                    }
                    channel.Writer.Complete();
                }
                catch (Exception ex)
                {
                    channel.Writer.Complete(ex);
                }
            }
        }

        async Task RunPostProcessing(IAsyncEnumerable<PatternResult> source, ExecutionPlan plan)
        {
            using var writer = new ResultWriter(repo.NewContext(), plan);
            using var pub = new Publisher(plan.Options.OutputDirectory);
            await foreach (var result in source.ConfigureAwait(false))
            {
                pub.Publish(result);
                await writer.Write(result).ConfigureAwait(false);
            }
            await writer.WriteEndDate(DateTimeOffset.Now).ConfigureAwait(false);
        }

        async Task NopPostProcessing(IAsyncEnumerable<PatternResult> source, ExecutionPlan plan)
        {
            await foreach (var _ in source.ConfigureAwait(false))
            {
            }
        }

        async Task<int> RunImagingPipeline(ImagingContext ctx, ExecutionPlan plan)
        {
            var nav = await Align(ctx, plan).ConfigureAwait(false);
            var pipeline = ImagingPipeline(ctx, nav, plan);
            foreach (var p in plan.ScanPoints)
            {
                var c = new Coordinate(p.X, p.Y, p.Dx, p.Dy, p.AddrOffsetX, p.AddrOffsetY);
                await pipeline.MoveTo(c, p.Pattern).ConfigureAwait(false);
            }
            return 0;
        }

        ImagingPipeline ImagingPipeline(ImagingContext ctx, IWaferNavigator nav, ExecutionPlan plan)
        {
            return new ImagingPipeline(ctx, nav, semClient, imagingService, plan.Environment, plan.Options);
        }

        async Task<IWaferNavigator> Align(ImagingPipeline pipeline, Pattern pattern)
        {
            var alignment = new WaferAlignment(imagingService, stageService, pipeline);
            var options = await alignment.Align(pattern).ConfigureAwait(false);
            await semClient.Settings.SetAngle(options.Angle);
            return stageService.WaferNavigator(options);
        }

        async Task<IWaferNavigator> Align(ImagingContext ctx, ExecutionPlan plan)
        {
            if (plan.Alignment.Count == 0)
            {
                return stageService.Navigator;
            }

            await imagingService.UseHigh().ConfigureAwait(false);
            var pipeline = ImagingPipeline(ctx, stageService.WaferNavigator(plan.Map), plan);
            var nav = await Align(pipeline, plan.Alignment[0]).ConfigureAwait(false);
            await imagingService.UseSEM().ConfigureAwait(false);
            nav = stageService.WaferNavigator(nav.Options);
            return await Align(ImagingPipeline(ctx, nav, plan), plan.Alignment[1]).ConfigureAwait(false);
        }

        async Task<IWaferNavigator> UnsafeAlign(ImagingContext ctx, ExecutionPlan plan)
        {
            await imagingService.UseHigh().ConfigureAwait(false);
            var pipeline = ImagingPipeline(ctx, stageService.WaferNavigator(plan.Map), plan);
            var nav = await Align(pipeline, plan.Alignment[0]).ConfigureAwait(false);
            await imagingService.UnsafeUseSEM().ConfigureAwait(false);
            nav = stageService.WaferNavigator(nav.Options);
            return await Align(ImagingPipeline(ctx, nav, plan), plan.Alignment[1]).ConfigureAwait(false);
        }

        public void OnJobCompleted()
        {
            isBusy = false;
        }

        static void FixHorizontalAlignment(List<Coordinate> coordinates)
        {
            if (coordinates.Count > 2)
            {
                if (coordinates[0].Y != coordinates[1].Y)
                {
                    (coordinates[0], coordinates[2]) = (coordinates[2], coordinates[0]);
                }
            }
        }

        static List<Pattern> OptimizeAlignment(List<Pattern> patterns)
        {
            if (patterns.Count < 2)
            {
                return patterns;
            }

            var coordinates = patterns[1].Coordinates;
            FixHorizontalAlignment(coordinates);

            var x0 = patterns[0].Coordinates[1].X;
            var d1 = Math.Abs(coordinates[0].X - x0);
            var d2 = Math.Abs(coordinates[1].X - x0);
            if (d2 < d1)
            {
                (coordinates[0], coordinates[1]) = (coordinates[1], coordinates[0]);
            }

            return patterns;
        }

        static Pattern Copy(Pattern p)
        {
            return new Pattern
            {
                Kind = p.Kind,
                Ordinal = p.Ordinal,
                Dx = p.Dx,
                Dy = p.Dy,
                AddrOffsetX = p.AddrOffsetX,
                AddrOffsetY = p.AddrOffsetY,
                Reference = p.Reference,
                Addressing = p.Addressing,
                CDPoints = p.CDPoints,
                UseAddressing = p.UseAddressing,
                UseDifferentialDrive = p.UseDifferentialDrive
            };
        }
    }
}
